//////////////////////////////////////////////
// main.cpp
//
//////////////////////////////////////////////

/// Includes ---------------------------------

// Local
#include "PastCode.h"

// nkGraphics
#include <NilkinsGraphics/Buffers/Buffer.h>
#include <NilkinsGraphics/Buffers/BufferManager.h>

#include <NilkinsGraphics/Compositors/Compositor.h>
#include <NilkinsGraphics/Compositors/CompositorManager.h>
#include <NilkinsGraphics/Compositors/CompositorNode.h>
#include <NilkinsGraphics/Compositors/TargetOperations.h>

#include <NilkinsGraphics/Passes/ClearTargetsPass.h>
#include <NilkinsGraphics/Passes/ComputePass.h>
#include <NilkinsGraphics/Passes/PostProcessPass.h>
#include <NilkinsGraphics/Passes/RenderScenePass.h>

#include <NilkinsGraphics/RenderContexts/RenderContext.h>
#include <NilkinsGraphics/RenderContexts/RenderContextDescriptor.h>
#include <NilkinsGraphics/RenderContexts/RenderContextManager.h>

/// Internals : Resources --------------------

void prepareBuffer ()
{
	// A buffer is a resource holding binary information
	// It has no real meaning required, unlike a texture
	// They can be used in compute, pixel, and so on, once rightly setup
	nkGraphics::Buffer* buffer = nkGraphics::BufferManager::getInstance()->createOrRetrieve("filterBuffer") ;

	// We will use this buffer in a compute pass, but also as a resource later to copy it to the screen
	buffer->prepareForComputeResourceUsage() ;
	buffer->prepareForShaderResourceUsage() ;

	// Now we can setup the buffer's size
	// We will consider it as an array of pixels of 4 bytes (R8G8B8A8 -> 32 bits)
	buffer->setElementByteSize(4) ;
	// Our image is 800x600 pixels
	buffer->setElementCount(800 * 600) ;

	// A buffer is a resource like any other
	buffer->load() ;
}

void prepareOffscreenTexture ()
{
	// This texture will be used as a render target to keep the render result available to the compute stage
	// Its setup will be a little more manual than simply loading a file
	nkGraphics::Texture* tex = nkGraphics::TextureManager::getInstance()->createOrRetrieve("sceneTarget") ;

	// First, set up its size to the window size
	tex->setWidth(800) ;
	tex->setHeight(600) ;
	// Then it requires its format
	tex->setTextureFormat(nkGraphics::R8G8B8A8_UNORM) ;
	// A texture that can be used as a render target needs to be marked as such
	tex->setRenderFlag(nkGraphics::TEX_RENDER_FLAG::RENDER_TARGET) ;

	// We can request the load, everything is setup
	tex->load() ;
}

/// Internals : Shaders ----------------------

void prepareFilterProgram ()
{
	nkGraphics::Program* filterProgram = nkGraphics::ProgramManager::getInstance()->createOrRetrieve("filterProgram") ;
	nkGraphics::ProgramSourcesHolder sources ;

	// This program will only use compute stage
	sources.setComputeMemory
	(
		R"eos(
			cbuffer passConstants
			{
				uint4 texInfos ;
			}

			RWStructuredBuffer<uint> bufOut : register(u0) ;

			Texture2D inTex : register(t0) ;

			[numthreads(32, 32, 1)]
			void main (int3 dispatchThreadID : SV_DispatchThreadID)
			{
				if (dispatchThreadID.x < texInfos.x && dispatchThreadID.y < texInfos.y)
				{
					float4 texCol = inTex.Load(uint3(dispatchThreadID.xy, 0)) ;

					const float3x3 sepiaMask = 
					{
						0.393, 0.349, 0.272,
						0.769, 0.686, 0.534,
						0.189, 0.168, 0.131
					} ;

					texCol.rgb = saturate(mul(texCol.rgb, sepiaMask)) ;

					uint texColPacked = 0 ;
					texColPacked += ((uint)(texCol.r * 255)) << 24 ;
					texColPacked += ((uint)(texCol.g * 255)) << 16 ;
					texColPacked += ((uint)(texCol.b * 255)) << 8 ;
					texColPacked += 0xFF ;

					int index = dispatchThreadID.x + dispatchThreadID.y * texInfos.x ;

					bufOut[index] = texColPacked ;
				}
			}
		)eos"
	) ;

	filterProgram->setFromMemory(sources) ;
	filterProgram->load() ;
}

void prepareFilterShader ()
{
	// Prepare the shader used to apply the filter
	nkGraphics::Shader* filterShader = nkGraphics::ShaderManager::getInstance()->createOrRetrieve("filterShader") ;
	nkGraphics::Program* filterProgram = nkGraphics::ProgramManager::getInstance()->get("filterProgram") ;

	filterShader->setProgram(filterProgram) ;

	// Constant buffer only needs the size of the target used
	nkGraphics::ConstantBuffer* cBuffer = filterShader->addConstantBuffer(0) ;

	// We will find it through our offscreen texture
	nkGraphics::Texture* tex = nkGraphics::TextureManager::getInstance()->get("sceneTarget") ;
	nkGraphics::ShaderPassMemorySlot* slot = cBuffer->addPassMemorySlot() ;
	slot->setAsTextureSize(tex) ;

	// Prepare for texture
	filterShader->addTexture(tex, 0) ;

	// Finally, we need to bind our buffer to the UAV slots
	// An UAV slot allows for writing into it
	nkGraphics::Buffer* buffer = nkGraphics::BufferManager::getInstance()->get("filterBuffer") ;
	filterShader->addUavBuffer(buffer, 0) ;

	// Finalize loading
	filterShader->load() ;
}

void prepareBufferCopyProgram ()
{
	// This program will copy data from the buffer filled from our compute program and paste in unto the screen
	nkGraphics::Program* bufferCopyProgram = nkGraphics::ProgramManager::getInstance()->createOrRetrieve("bufferCopyProgram") ;
	nkGraphics::ProgramSourcesHolder sources ;

	sources.setVertexMemory
	(
		R"eos(
			struct VertexInput
			{
				float4 position : POSITION ;
				float2 uvs : TEXCOORD0 ;
			} ;

			struct PixelInput
			{
				float4 position : SV_POSITION ;
				float2 uvs : TEXCOORD0 ;
				uint4 texInfos : TEXINFOS ;
			} ;

			cbuffer passConstants
			{
				uint4 texInfos ;
			}

			PixelInput main (VertexInput input)
			{
				PixelInput result ;

				result.position = input.position ;
				result.uvs = input.uvs ;
				result.texInfos = texInfos ;

				return result ;
			}
		)eos"
	) ;

	sources.setPixelMemory
	(
		R"eos(
			struct PixelInput
			{
				float4 position : SV_POSITION ;
				float2 uvs : TEXCOORD0 ;
				uint4 texInfos : TEXINFOS ;
			} ;

			StructuredBuffer<uint> inputBuf : register(t0) ;

			float4 main (PixelInput input) : SV_TARGET
			{
				uint2 index = uint2(input.uvs.x * input.texInfos.x, input.uvs.y * input.texInfos.y) ;
				uint texColPacked = inputBuf[index.x + index.y * input.texInfos.x] ;

				uint texColR = (texColPacked & 0xFF000000) >> 24 ;
				uint texColG = (texColPacked & 0x00FF0000) >> 16 ;
				uint texColB = (texColPacked & 0x0000FF00) >> 8 ;
				float3 texCol = float3(texColR / 255.0, texColG / 255.0, texColB / 255.0) ;

				return float4(texCol, 1.0) ;
			}
		)eos"
	) ;

	bufferCopyProgram->setFromMemory(sources) ;
	bufferCopyProgram->load() ;
}

void prepareBufferCopyShader ()
{
	// Prepare the shader used to apply the filter
	nkGraphics::Shader* bufferCopyShader = nkGraphics::ShaderManager::getInstance()->createOrRetrieve("bufferCopyShader") ;
	nkGraphics::Program* bufferCopyProgram = nkGraphics::ProgramManager::getInstance()->get("bufferCopyProgram") ;

	bufferCopyShader->setProgram(bufferCopyProgram) ;

	// Constant buffer only needs the size of the target used
	nkGraphics::ConstantBuffer* cBuffer = bufferCopyShader->addConstantBuffer(0) ;

	// Target will drive this
	nkGraphics::ShaderPassMemorySlot* slot = cBuffer->addPassMemorySlot() ;
	slot->setAsTargetSize() ;

	// Then we need to bind our buffer to read it from
	// This time, API-wise, it is considered to be a texture, as we will only read from it
	nkGraphics::Buffer* buffer = nkGraphics::BufferManager::getInstance()->get("filterBuffer") ;
	bufferCopyShader->addTexture(buffer, 0) ;

	// Finalize loading
	bufferCopyShader->load() ;
}

/// Internals : Compositor -------------------

nkGraphics::Compositor* prepareCompositor ()
{
	// Prepare the shaders passes will require
	nkGraphics::Shader* envShader = nkGraphics::ShaderManager::getInstance()->get("envShader") ;
	nkGraphics::Shader* filterShader = nkGraphics::ShaderManager::getInstance()->get("filterShader") ;
	nkGraphics::Shader* bufferCopyShader = nkGraphics::ShaderManager::getInstance()->get("bufferCopyShader") ;

	// Prepare the textures target operations will require
	nkGraphics::Texture* sceneTarget = nkGraphics::TextureManager::getInstance()->get("sceneTarget") ;

	// Get the compositor
	nkGraphics::Compositor* compositor = nkGraphics::CompositorManager::getInstance()->createOrRetrieve("compositor") ;
	nkGraphics::CompositorNode* node = compositor->addNode() ;

	// This compositor will require more magic than the last one
	// We will :
	//     - First render into our offscreen texture [TargetOperations 0]
	//         + Clear the target                [ClearTargetsPass]
	//         + Render the sphere               [RenderScenePass]
	//         + Render the background           [PostProcessPass]
	//         + Filter the result               [ComputePass]
	//     - Render to the context's surface         [TargetOperations 1]
	//         + Copy the buffer to the texture  [PostProcessPass]
	// Using this composition pattern, we will be able to reach the result we want : a sepia filtered image rendered
	// First operation will render offscreen, but still use the context's depth buffer (no need for another)
	nkGraphics::TargetOperations* targetOp0 = node->addOperations() ;
	targetOp0->addColorTarget(sceneTarget) ;
	targetOp0->setToChainDepthBuffer(true) ;

	// Unroll our passes
	nkGraphics::ClearTargetsPass* clearPass = targetOp0->addClearTargetsPass() ;
	nkGraphics::RenderScenePass* scenePass = targetOp0->addRenderScenePass() ;

	nkGraphics::PostProcessPass* postProcessPass = targetOp0->addPostProcessPass() ;
	postProcessPass->setBackProcess(true) ;
	postProcessPass->setShader(envShader) ;

	// Addition this time : the compute pass
	nkGraphics::ComputePass* computePass = targetOp0->addComputePass() ;
	computePass->setShader(filterShader) ;
	// We know threads spawned will be 32x32, and each thread will filter one pixel
	// Our image being 800x600, we will require at least (2 * 32)5 groups on width, and (19 * 32) groups on height
	computePass->setX(25) ;
	computePass->setY(19) ;

	// Prepare next target operations
	// This one will only be responsible for copying the buffer to the rendering surface
	nkGraphics::TargetOperations* targetOp1 = node->addOperations() ;
	targetOp1->setToBackBuffer(true) ;
	
	nkGraphics::PostProcessPass* copyPass = targetOp1->addPostProcessPass() ;
	copyPass->setBackProcess(false) ;
	copyPass->setShader(bufferCopyShader) ;

	return compositor ;
}

/// Function ---------------------------------

int main ()
{
	// Prepare for logging
	std::unique_ptr<nkLog::Logger> logger = std::make_unique<nkLog::ConsoleLogger>() ;
	nkGraphics::LogManager::getInstance()->setReceiver(logger.get()) ;

	// For easiness
	nkResources::ResourceManager::getInstance()->setWorkingPath("Data") ;

	// Initialize and create context with window
	if (!nkGraphics::System::getInstance()->initialize())
		return -1 ;

	nkGraphics::RenderContext* context = nkGraphics::RenderContextManager::getInstance()->createRenderContext(nkGraphics::RenderContextDescriptor(800, 600, false, true)) ;

	baseInit() ;

	// Prepare the texture we will use
	prepareBuffer() ;
	prepareOffscreenTexture() ;

	// Filtering shaders
	prepareFilterProgram() ;
	prepareFilterShader() ;
	prepareBufferCopyProgram() ;
	prepareBufferCopyShader() ;

	// Prepare the composition once everything is ready
	nkGraphics::Compositor* compositor = prepareCompositor() ;
	
	// Use the compositor for the context we just created
	context->setCompositor(compositor) ;

	// And trigger the rendering
	renderLoop(context) ;

	// Clean exit
	nkGraphics::System::getInstance()->kill() ;

	return 0 ;
}